{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 16\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "EPOCHS = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline=transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.1307,),(0.3081,))\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "train_set=datasets.MNIST(\"data\",train=True ,download=True,transform=pipeline)\n",
    "test_set=datasets.MNIST(\"data\",train=False ,download=True,transform=pipeline)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader=DataLoader(train_set,batch_size=BATCH_SIZE,shuffle=True)\n",
    "test_loader=DataLoader(test_set,batch_size=BATCH_SIZE,shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(28, 28, 1)\n"
     ]
    }
   ],
   "source": [
    "with open (\"./data/MNIST/raw/train-images-idx3-ubyte\",\"rb\") as f:\n",
    "    file =f.read()\n",
    "    image1=[int(str(item).encode('ascii'),16) for item in file[16 :16+784]]\n",
    "    # print(image1)\n",
    "    import cv2\n",
    "    import numpy as np\n",
    "    image1_np=np.array(image1,dtype=np.uint8).reshape(28,28,1)\n",
    "    print(image1_np.shape)\n",
    "    cv2.imwrite(\"digit1.jpg\",image1_np)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Digit(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 10, 5)\n",
    "        self.conv2 = nn.Conv2d(10, 20, 3)  # kernel is 3*3\n",
    "        self.fc1 = nn.Linear(20*10*10, 500)\n",
    "        self.fc2 = nn.Linear(500, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        input_size = x.size(0)\n",
    "        # batch_size * 1 * 28 * 28( 0 means the first one -- batch_size )\n",
    "\n",
    "        x = self.conv1(x)\n",
    "        # input : batchsize * 1 * 28 *28 ,output : batchsize * 10 * 24 *24\n",
    "        x = F.relu(x)\n",
    "\n",
    "        x = F.max_pool2d(x, 2, 2)\n",
    "        # input : batchsize * 10 * 24 * 24 ,output : batchsize * 10 * 12 * 12\n",
    "\n",
    "        x = self.conv2(x)\n",
    "        # input : batchsize * 10 * 12 * 12 ,output : batchsize * 20 * 10 * 10\n",
    "        x=F.relu(x)\n",
    "\n",
    "        x=x.view(input_size,-1)\n",
    "        #laping , -1 is calculate the dimision : 20 * 10 * 10 = 2000\n",
    "\n",
    "        x=self.fc1(x)\n",
    "        #input : batchsize * 2000 output : batchsize * 500\n",
    "        x=F.relu(x)\n",
    "\n",
    "        x=self.fc2(x)\n",
    "        #input : batchsize * 500 output : batchsize * 10\n",
    "\n",
    "        output =F.log_softmax(x,dim=1)\n",
    "        #calculate the % of each number\n",
    "\n",
    "        return output\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Digit().to(DEVICE)\n",
    "optimizer = optim.Adam(model.parameters())\n",
    "# update model canshu\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, device, train_loader, optimizer, epoch):\n",
    "    model.train()\n",
    "    for batch_index, (data, target) in enumerate(train_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = F.cross_entropy(output, target)\n",
    "        # _,pred = output.max(1, keepdim=True)\n",
    "        # or pred=output.max(1)[1]\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        if batch_index % 100 == 0:\n",
    "            print(\"Train epoch : {} \\t loss : {:.6f}\".format(\n",
    "                epoch, loss.item()))\n",
    "            # print(target)\n",
    "            # print(pred)\n",
    "\n",
    "\n",
    "# train_model(model, DEVICE, train_loader, optimizer, EPOCHS)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_model(model,device,test_loader):\n",
    "    model.eval()\n",
    "    correct=0.0\n",
    "    test_loss=0.0\n",
    "    with torch.no_grad():# 不计算梯度，不会反向传播\n",
    "        for batch_index,(data,target) in enumerate(test_loader):\n",
    "            data,target = data.to(device),target.to(device)\n",
    "            output = model(data)\n",
    "            test_loss+=F.cross_entropy(output,target).item()\n",
    "            pred=output.max(1,keepdim=True)[1] # [1] is pred.indice\n",
    "            # pred = torch.max(output,dim=1)\n",
    "            # pred =output.argmax(dim=1)\n",
    "            fig = plt.figure()\n",
    "            plt.tight_layout()\n",
    "            x=data[2]\n",
    "            x = x.swapaxes(0, 1)\n",
    "            x = x.swapaxes(1, 2)\n",
    "            plt.imshow(x.squeeze())\n",
    "            plt.title(\"Ground Truth: {}\".format(output.max(1,keepdim=True)[1][2]))\n",
    "            plt.xticks([])\n",
    "            plt.yticks([])\n",
    "            plt.show()\n",
    "\n",
    "            correct+=pred.eq(target.view_as(pred)).sum().item()\n",
    "        test_loss/=len(test_loader.dataset)\n",
    "        print(\"Test --ave loss : {: .4f} , Accuracy : {:.3f}%\\n\".format(\n",
    "            test_loss ,100*correct/len(test_loader.dataset)\n",
    "        ))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for epoch in range(1,EPOCHS+1)\n",
    "train_model(model,DEVICE,train_loader,optimizer,0)\n",
    "test_model(model,DEVICE,test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
